#include <iostream>
#include <immintrin.h>
#include <windows.h>

#define VECTOR_SIZE 8 

void gemm(float *A, float *B, float *C, int m, int n, int k){
    int p1 = m*k;
    int all_size = p1 + n*k;

    float *all = new float[all_size];
    for (int i = 0; i < p1; i++){
        all[i] = A[i];
    }

    int idx = p1;
    for (int i = 0; i < n; i++){
        for (int j = 0; j < k; j++){
            all[idx] = B[j*n + i];
            idx++;
        }
    }
    int h = m + n;
    idx = 0;
    
    float * temp = new float[k];
    for (int i = 0; i < m; i++){
        for (int j = m; j < h; j++){
            float *s1 = &all[i*k];
            float *s2 = &all[j*k];
            register float sum = 0.0f;
            int c;
            for (c = 0; c <= k - VECTOR_SIZE; c += VECTOR_SIZE) {
                __m256 av = _mm256_loadu_ps(&s1[c]); 
                __m256 bv = _mm256_loadu_ps(&s2[c]); 
                __m256 rv = _mm256_mul_ps(av, bv);  
                _mm256_storeu_ps(&temp[c], rv);   
            }

            for (; c < k; c++) {
                temp[c] = s1[c] * s2[c];
            }

            for (int w = 0; w < k; w++){
                sum += temp[w];
            }
            C[idx] = sum;
            idx++;
        }
    }
    delete[] temp;
    delete[] all;

}

int main() {
    int m = 512, n = 512, k = 512;
    float *A = new float[m*k];
    float *B = new float[k*n];
    float *C = new float[m*n];

    for (int i = 0; i < m*k; i++){
        A[i] = 1.0f;
    }

    for (int i = 0; i < k*n; i++){
        B[i] = 1.0f;
    }

    for (int i = 0; i < m*n; i++){
        C[i] = 0.0f;
    }

    DWORD start_t, end_t;

    start_t = GetTickCount();
    
    gemm(A, B, C, m , n, k);

    end_t = GetTickCount();

    double time = (end_t - start_t)/1e+3;

    std::cout << "time: " << time << "s" << std::endl;

    double flops = 2.0 * m * n * k * 1e-9;
    std::cout << "flops: " << flops/time << "Gflops" << std::endl;

    delete[] A;
    delete[] B;
    delete[] C;
    return 0;
}
